import json
import random
import time
import argparse
import logging
from playwright.sync_api import sync_playwright, TimeoutError

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

# Configuration
CONFIG = {
    "browser_ws_endpoint": "http://localhost:9222",
    "allowed_domains": ["new.oaifree.com", "chatgpt.com"],
    "min_delay": 25,
    "max_delay": 35,
    "prompt_styles": {
        "default": {},
        "photography": {"post_prompt": " 3D渲染，真实场景，可爱风格，方形宽高比"},
        "3d_render": {"pre_prompt": " 3D render of an adorable ，方形宽高比，拟人"},
        "cute_style": {"pre_prompt": "可爱风格，电影3D渲染，方形宽高比，拟人"},
        "universal": {"pre_prompt": "电影特效，方形宽高比，拟人，"},
        "example": {"pre_prompt": "电影特效，拟人，", "post_prompt": "，方形宽高比"}
    }
}

def load_json(file_path):
    try:
        with open(file_path, 'r', encoding='utf-8') as file:
            data = json.load(file)
        return [{"scene": i+1, "description": scene["description"]} for i, scene in enumerate(data["storyboard"])]
    except FileNotFoundError:
        logging.error(f"JSON file not found: {file_path}")
        raise
    except json.JSONDecodeError:
        logging.error(f"Invalid JSON in file: {file_path}")
        raise

def load_txt(file_path):
    try:
        with open(file_path, 'r', encoding='utf-8') as file:
            lines = file.readlines()
        return [{"scene": i+1, "description": line.strip()} for i, line in enumerate(lines) if line.strip()]
    except FileNotFoundError:
        logging.error(f"TXT file not found: {file_path}")
        raise

def load_file(file_path):
    if file_path.lower().endswith('.json'):
        return load_json(file_path)
    elif file_path.lower().endswith('.txt'):
        return load_txt(file_path)
    else:
        logging.error(f"Unsupported file type: {file_path}")
        raise ValueError("Unsupported file type. Please use .json or .txt files.")

def generate_message(description, style):
    pre_prompt = style.get("pre_prompt", "")
    post_prompt = style.get("post_prompt", "")
    message = f"{pre_prompt} {description} {post_prompt}".strip()
    return message

def wait_for_send_button(page, timeout=180):
    start_time = time.time()
    while time.time() - start_time < timeout:
        send_button = page.get_by_test_id("send-button")
        if send_button.is_visible():
            return send_button
        time.sleep(10)
        logging.info("Waiting for send button to become visible...")
    
    logging.error(f"Send button not found after {timeout} seconds. Stopping the program.")
    raise TimeoutError("Send button not found")

def send_message(page, message, min_delay, max_delay):
    try:
        # Find the input field and send the message
        input_field = page.locator('#prompt-textarea')
        input_field.fill(message)
        logging.info(f"Message filled: {message[:50]}...")  # Log first 50 characters
        
        # Find and click the send button
        send_button = wait_for_send_button(page)
        send_button.click()
        logging.info("Message sent successfully")
        
        # Wait for a random time before next message
        delay = random.uniform(min_delay, max_delay)
        time.sleep(delay)
    except TimeoutError:
        logging.error("Timeout error: Send button not found or not clickable")
        raise
    except Exception as e:
        logging.error(f"Error sending message: {str(e)}")
        raise

def main(file_path, prompt_style, start_scene, specific_scenes):
    data = load_file(file_path)
    style = CONFIG["prompt_styles"].get(prompt_style, CONFIG["prompt_styles"]["default"])
    browser = None
    
    try:
        with sync_playwright() as p:
            try:
                browser = p.chromium.connect_over_cdp(CONFIG['browser_ws_endpoint'])
                logging.info(f"Successfully connected to browser at {CONFIG['browser_ws_endpoint']}")
            except Exception as e:
                logging.error(f"Failed to connect to browser: {str(e)}")
                return

            try:
                context = browser.contexts[0]
                pages = context.pages
                page = pages[0]
                logging.info(f"Page URL: {page.url}")
                logging.info(f"Page title: {page.title()}")
            except IndexError:
                logging.error("No pages found in the browser context")
                return
            
            # Check if the current page matches any of the allowed domains
            current_url = page.url
            if not any(domain in current_url for domain in CONFIG['allowed_domains']):
                logging.warning(f"Current page {current_url} is not in the list of allowed domains.")
                return
            
            total_scenes = len(data)
            scenes_to_run = specific_scenes if specific_scenes else range(start_scene, total_scenes + 1)
            
            for i in scenes_to_run:
                if i > total_scenes:
                    logging.warning(f"Scene {i} does not exist. Skipping.")
                    continue
                scene = data[i-1]
                description = scene["description"]
                message = generate_message(description, style)
                logging.info(f"Processing scene {i}/{total_scenes}")
                try:
                    send_message(page, message, CONFIG['min_delay'], CONFIG['max_delay'])
                except TimeoutError:
                    logging.error("Stopping the program due to timeout.")
                    break
                except Exception as e:
                    logging.error(f"Failed to send message for scene {i}: {str(e)}")
                    continue  # Continue with the next scene
            
            logging.info("Script completed successfully")
    except KeyboardInterrupt:
        logging.info("Script interrupted by user")
    except Exception as e:
        logging.error(f"An unexpected error occurred: {str(e)}")
    finally:
        if browser:
            try:
                browser.close()
                logging.info("Browser closed successfully")
            except Exception as e:
                logging.error(f"Error closing browser: {str(e)}")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Automated message sender script")
    parser.add_argument("file_path", help="Path to the JSON or TXT file containing the storyboard")
    parser.add_argument("--prompt_style", default="default", choices=CONFIG["prompt_styles"].keys(), 
                        help="Style of prompt to use (default: default)")
    parser.add_argument("--start_scene", type=int, default=1, help="Scene number to start from (default: 1)")
    parser.add_argument("--scenes", type=str, help="Comma-separated list of specific scenes to run (e.g., '3,4,6,8')")
    
    args = parser.parse_args()
    
    specific_scenes = [int(s) for s in args.scenes.split(',')] if args.scenes else None
    
    main(args.file_path, args.prompt_style, args.start_scene, specific_scenes)